import sys
import os
from typing import Optional, List
from contextlib import asynccontextmanager
import warnings

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 忽略模型缓存符号链接警告
warnings.filterwarnings("ignore", message="Failed to create symbolic link")

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Response, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

# 从正确的路径导入模块
from backend.cache import RedisCache
from backend.danwen.dquestion_milvus import SQLiteDatabaseRAG

# 加载环境变量
load_dotenv()

# 定义lifespan上下文管理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    应用程序生命周期管理
    在应用启动时初始化数据库连接和索引，应用结束时清理资源
    """
    # 在应用启动时执行
    global vn
    await vn.create_index()
    yield
    # 在应用关闭时执行
    # 这里可以添加清理代码

# 初始化缓存和数据库RAG实例
cache = RedisCache()
vn = SQLiteDatabaseRAG()

# 创建FastAPI应用
app = FastAPI(title="Text2SQL UI", lifespan=lifespan)

# 获取静态文件的绝对路径
static_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
assets_dir = os.path.join(static_dir,"dist", "assets")

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory=static_dir), name="static")
# 挂载assets目录到根路径下，这样可以直接通过/assets/访问
app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")


# # 添加单个静态文件的路由
# @app.get("/vanna.svg")
# async def get_vanna_svg():
#     """返回Vanna SVG图标"""
#     vanna_svg_path = os.path.join(static_dir, "vanna.svg")
#     with open(vanna_svg_path, "rb") as f:
#         return Response(content=f.read(), media_type="image/svg+xml")


# 配置CORS（跨域资源共享）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,  # 允许携带凭证
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头部
)

# 定义请求模型
class TrainingData(BaseModel):
    question: Optional[str] = None  # 问题文本
    sql: Optional[str] = None  # SQL查询语句
    ddl: Optional[str] = None  # 数据定义语言
    documentation: Optional[str] = None  # 文档说明


class TrainingDataRemoval(BaseModel):
    id: str  # 训练数据ID


# 工具函数：检查缓存中是否存在指定字段
async def check_cache(id: str, fields: List[str]):
    """检查缓存中是否存在指定字段，如果不存在则抛出异常"""
    if not id:
        raise HTTPException(status_code=400, detail="No id provided")

    for field in fields:
        if cache.get(id=id, field=field) is None:
            raise HTTPException(status_code=404, detail=f"No {field} found")

    return {field: cache.get(id=id, field=field) for field in fields}


# 添加根路径重定向到静态index.html
@app.get("/")
async def redirect_to_index():
    """重定向到静态首页"""
    return Response(
        status_code=307,
        headers={"Location": "/static/dist/index.html"}
    )


# 生成问题列表接口
@app.get("/api/v0/generate_questions", summary="生成问题列表")
async def generate_questions():
    """生成可以用来提问的问题列表"""
    return {
        "type": "question_list",
        "questions": vn.generate_questions(),
        "header": "Here are some questions you can ask:"
    }


# 根据问题生成SQL接口
@app.get("/api/v0/generate_sql", summary="根据问题生成SQL")
async def generate_sql(question: str = Query(None)):
    """根据自然语言问题生成SQL查询"""
    if not question:
        raise HTTPException(status_code=400, detail="No question provided")

    # 为问题生成唯一ID
    id = cache.generate_id(question=question)
    # 生成对应的SQL
    sql = vn.generate_sql(question=question, allow_llm_to_see_data=True)

    # 存入缓存
    cache.set(id=id, field='question', value=question)
    cache.set(id=id, field='sql', value=sql)

    return {
        "type": "sql",
        "id": id,
        "text": sql,
    }


# 执行SQL查询接口
@app.get("/api/v0/run_sql", summary="执行SQL查询")
async def run_sql(id: str):
    """执行缓存中的SQL查询并返回结果"""
    # 检查缓存中是否有SQL
    data = await check_cache(id, ['sql'])

    try:
        # 执行SQL查询
        df = vn.run_sql(sql=data['sql'])

        # 将结果存入缓存
        cache.set(id=id, field='df', value=df)

        return {
            "type": "df",
            "id": id,
            "df": df.head(10).to_json(orient='records'),  # 返回前10行
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# 下载CSV文件接口
@app.get("/api/v0/download_csv", summary="下载CSV文件")
async def download_csv(id: str):
    """将查询结果下载为CSV文件"""
    # 从缓存获取数据框
    data = await check_cache(id, ['df'])

    # 转换为CSV格式
    csv = data['df'].to_csv()

    headers = {
        "Content-Disposition": f"attachment; filename={id}.csv"
    }

    return Response(content=csv, media_type="text/csv", headers=headers)


# 生成Plotly图表接口
@app.get("/api/v0/generate_plotly_figure", summary="生成Plotly图表")
async def generate_plotly_figure(id: str):
    """根据查询结果生成Plotly可视化图表"""
    # 检查缓存中的必要数据
    data = await check_cache(id, ['df', 'question', 'sql'])

    try:
        # 获取数据框的元数据
        df_metadata = f"Running df.dtypes gives:\n {data['df'].dtypes}"

        # 生成Plotly可视化代码
        code = vn.generate_plotly_code(
            question=data['question'],
            sql=data['sql'],
            df_metadata=df_metadata
        )

        # 执行代码获取图表
        fig = vn.get_plotly_figure(plotly_code=code, df=data['df'], dark_mode=False)
        fig_json = fig.to_json()

        # 保存图表JSON到缓存
        cache.set(id=id, field='fig_json', value=fig_json)

        return {
            "type": "plotly_figure",
            "id": id,
            "fig": fig_json,
        }
    except Exception as e:
        # 打印堆栈跟踪便于调试
        import traceback
        traceback.print_exc()

        raise HTTPException(status_code=400, detail=str(e))


# 获取训练数据接口
@app.get("/api/v0/get_training_data", summary="获取训练数据")
async def get_training_data():
    """获取当前系统中的训练数据"""
    df = vn.get_training_data()

    return {
        "type": "df",
        "id": "training_data",
        "df": df.head(25).to_json(orient='records'),  # 返回前25行
    }


# 删除训练数据接口
@app.post("/api/v0/remove_training_data", summary="删除训练数据")
async def remove_training_data(data: TrainingDataRemoval):
    """删除指定ID的训练数据"""
    if vn.remove_training_data(id=data.id):
        return {"success": True}

    raise HTTPException(status_code=400, detail="Couldn't remove training data")


# 添加训练数据接口
@app.post("/api/v0/train", summary="添加训练数据")
async def add_training_data(data: TrainingData):
    """添加新的训练数据到系统中"""
    try:
        id = vn.train(
            question=data.question,
            sql=data.sql,
            ddl=data.ddl,
            documentation=data.documentation
        )

        return {"id": id}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# 生成后续问题接口
@app.get("/api/v0/generate_followup_questions", summary="生成后续问题")
async def generate_followup_questions(id: str):
    """根据当前问题和结果生成后续可能的问题"""
    # 检查缓存中的必要数据
    data = await check_cache(id, ['df', 'question', 'sql'])

    # 生成后续问题
    followup_questions = vn.generate_followup_questions(
        question=data['question'],
        sql=data['sql'],
        df=data['df']
    )

    # 将后续问题存入缓存
    cache.set(id=id, field='followup_questions', value=followup_questions)

    return {
        "type": "question_list",
        "id": id,
        "questions": followup_questions,
        "header": "Here are some followup questions you can ask:"
    }


# 加载本地问题接口
@app.get("/api/v0/load_question", summary="加载本地问题")
async def load_question(id: str):
    """从缓存中加载完整的问题和结果数据"""
    # 从缓存中获取所有相关数据
    data = await check_cache(id, ['question', 'sql', 'df', 'fig_json', 'followup_questions'])

    try:
        return {
            "type": "question_cache",
            "id": id,
            "question": data['question'],
            "sql": data['sql'],
            "df": data['df'].head(10).to_json(orient='records'),
            "fig": data['fig_json'],
            "followup_questions": data['followup_questions'],
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# 获取问题历史记录接口
@app.get("/api/v0/get_question_history", summary="获取问题历史记录")
async def get_question_history():
    """获取用户之前提问的历史记录"""
    return {
        "type": "question_history",
        "questions": cache.get_all(field_list=['question'])
    }


# 主程序入口
if __name__ == "__main__":
    import uvicorn

    print("服务启动: http://127.0.0.1:8000")
    print("API文档: http://127.0.0.1:8000/docs")
    # 启动FastAPI服务，监听所有网络接口的8000端口，开启热重载
    uvicorn.run("fastapi_app:app", host="0.0.0.0", port=8000, reload=True)
